#encoding=utf-8
from openai import OpenAI
from api.global_config import LLAMA3_CONFIG
from util.post_processing import extract_sentences, get_stripped_sentences, find_word_in_string
import string


class LLAMA3:
    def __init__(self, config: LLAMA3_CONFIG):
        self.config = config

    def make_requests(self, prompt, isStream=False):
        # print(f"Calling LLAMA3 API with config: {self.config.__dict__}")
        client = OpenAI(
            base_url="https://integrate.api.nvidia.com/v1",
            api_key=self.config.api_key,
        )

        response = client.chat.completions.create(
            model=self.config.engine,
            messages=[
                {"role": "user", "content": prompt}
            ],
            temperature=self.config.temperature,
            top_p=self.config.top_p,
            max_tokens = self.config.max_tokens,
            stream=isStream
        )
        if isStream: # 仅在测试中使用，打印后没有后续流程
            for chunk in response:
                if chunk.choices[0].delta.content is not None:
                    print(chunk.choices[0].delta.content, end="")
        else:
            return response

    def post_processing(self, response):
        if response is None or response.choices[0].finish_reason=="length":
            return []
        raw_instructions = get_stripped_sentences(extract_sentences(response.choices[0].message.content))
        instructions = []
        for inst in raw_instructions:
            if inst == "":
                continue
            # filter out too short or too long instructions
            if len(inst) <= 3 or len(inst) > 150:
                continue
            # filter based on keywords that are not suitable for language models.
            if any(find_word_in_string(word, inst) for word in
                   ["image", "images", "graph", "graphs", "picture", "pictures", "file", "files", "map", "maps", "draw",
                    "plot", "go to"]):
                continue
            # We found that the model tends to add "write a program" to some existing instructions, which lead to a lot of such instructions.
            # And it's a bit confusing whether the model need to write a program or directly output the result.
            # Here we filter them out.
            # Note this is not a comprehensive filtering for all programming instructions.
            if inst.startswith("Write a program"):
                continue
            # filter those starting with punctuation
            if inst[0] in string.punctuation:
                continue
            # filter those starting with numbers
            if inst[0].isnumeric():
                continue
            instructions.append(inst)

        return instructions, response

    def get_raw_output(self, prompt):
        return self.make_requests(prompt).choices[0].message.content

    def get_instructions(self, prompt):
        return self.post_processing(self.make_requests(prompt))


if __name__ == "__main__":
    # NVIDIA_API_KEY = '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'
    # LLAMA_ENGINE = "meta/llama3-70b-instruct"
    prompt="Come up with a series of tasks with the same format:1.What's your name? 2.How's the weather? 3.How old are you? 4."
    prompt_chinese = 'Present a series of questions following the example below, requesting their content to be relevant to the example, but do not repeat the example:\n1. 如何使用信息技术为思政教育赋能？\n2. 请详细说明劳模精神、劳动精神、工匠精神的内涵和作用分别是什么，以及如何培育这些精神\n3. 红蓝融合中，理论和实践哪一个更重要？\n4. 见义勇为是否值得提倡？\n5. 信息时代，如何激发创新活力，产出创新成果？\n6. 请列举几个知行合一的名人事例\n7. 领导让我作一周工作的汇报总结，但我这周没有做实质性的工作，我该怎么办？\n8. 为什么要讲红蓝融合，其内涵是什么？\n9.'

    prompt = '''
    You are an advanced assistant to answer the question with the help of given relevant text.\nWhen user give you some texts, you may find that some of them is not relevant to thequestion given by user, now you need to identify the correspondence and overlook thenot relevant.\nMay be you have had the knowledge about the question, you better answer with the text,which is expected by user.\nYou can first determine the topic of question ,them use it as direction to give the answer.\nYou do not need to return the topic or quote!\nPlease answer this question in Chinese:\n\ntransformer是什么？\nYou can use following text as auxiliary information:\n\n 这张图片展示的是一个Transformer模型的结构图。Transformer是一种深度学习模型，广泛应用于自然语言处理。\n\n左边是“编码器”（Encoder）部分，包含了多个相同的层，每一层都有两个子层。第一个子层是“多头注意力”（Multi-Head Attention）机制， 第二个是简单的前馈神经网络（Feed Forward）。每个子层的输出都会通过一个“加法及规范化”（Add & Norm）步骤。此外，编码器的输入先通过一个“输入嵌入”（Input Embedding）层，然后加上“位置编码（Positional Encoding）。\n\n右边是“解码器”（Decoder）部分，也包含了多个相同的层，每一层有三个子层。第一个子层是“遮蔽的多头注意力”（Masked Multi-Head Attention）机制，用于防止位置后后的信息流入。第二个子层是一个普通的“多头注意力”机制，与编码器的最后一个层的输出相连。第三个子层是前馈神经网络。解码器的输入首先通过一个“输出嵌入”（Output Embedding）层，并进行右移，，后同样加上“位置编码”。\n\n最上方是该模型输出的步骤，包括一个线性层（Linear）和一个Softmax层，它们将解码器的输出转换为最终概率分布——模型的最终输出。\n\n这个图通常用于解释Transformer r模型的结构和用于序列到序列的学习任务，比如机器翻译。\n\nTo the best of our knowledge, however, the Transformer is the first transduction model relying\nentirely on self-attention to compute representations of its input and output without using sequence-\naligned RNNs or convolution. In the following sections, we will describe the Transformer, motivate\nself-attention and discuss its advantages over models such as [17, 18] and [9].\n3\nModel Architecture\nMost competitive neural sequence transduction models have an encoder-decoder structure [5, 2, 35].\n\nattention, replacing the recurrent layers most commonly used in encoder-decoder architectures with\nmulti-headed self-attention.\nFor translation tasks, the Transformer can be trained significantly faster than architectures based\non recurrent or convolutional layers. On both WMT 2014 English-to-German and WMT 2014\nEnglish-to-French translation tasks, we achieve a new state of the art. In the former task our best\nmodel outperforms even all previously reported ensembles.\n\nFigure 1: The Transformer - model architecture.\nThe Transformer follows this overall architecture using stacked self-attention and point-wise, fully\nconnected layers for both the encoder and decoder, shown in the left and right halves of Figure 1,\nrespectively.\n3.1\nEncoder and Decoder Stacks\nEncoder:\nThe encoder is composed of a stack of N = 6 identical layers. Each layer has two\nsub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position-\n\nThe Transformer allows for significantly more parallelization and can reach a new state of the art in\ntranslation quality after being trained for as little as twelve hours on eight P100 GPUs.\n2\nBackground\nThe goal of reducing sequential computation also forms the foundation of the Extended Neural GPU\n[16], ByteNet [18] and ConvS2S [9], all of which use convolutional neural networks as basic building\n\ntion models in various tasks, allowing modeling of dependencies without regard to their distance in\nthe input or output sequences [2, 19]. In all but a few cases [27], however, such attention mechanisms\nare used in conjunction with a recurrent network.\nIn this work we propose the Transformer, a model architecture eschewing recurrence and instead\nrelying entirely on an attention mechanism to draw global dependencies between input and output.\n\nprisingly well, yielding better results than all previously reported models with the exception of the\nRecurrent Neural Network Grammar [8].\nIn contrast to RNN sequence-to-sequence models [37], the Transformer outperforms the Berkeley-\nParser [29] even when training only on the WSJ training set of 40K sentences.\n7\nConclusion\nIn this work, we presented the Transformer, the first sequence transduction model based entirely on\n\nŁukasz Kaiser∗\nGoogle Brain\nlukaszkaiser@google.com\nIllia Polosukhin∗‡\nillia.polosukhin@gmail.com\nAbstract\nThe dominant sequence transduction models are based on complex recurrent or\nconvolutional neural networks that include an encoder and a decoder. The best\nperforming models also connect the encoder and decoder through an attention\nmechanism. We propose a new simple network architecture, the Transformer,\nbased solely on attention mechanisms, dispensing with recurrence and convolutions\n\nblock, computing hidden representations in parallel for all input and output positions. In these models,\nthe number of operations required to relate signals from two arbitrary input or output positions grows\nin the distance between positions, linearly for ConvS2S and logarithmically for ByteNet. This makes\nit more difficult to learn dependencies between distant positions [12]. In the Transformer this is\n\nTable 2: The Transformer achieves better BLEU scores than previous state-of-the-art models on the\nEnglish-to-German and English-to-French newstest2014 tests at a fraction of the training cost.\nModel\nBLEU\nTraining Cost (FLOPs)\nEN-DE\nEN-FR\nEN-DE\nEN-FR\nByteNet [18]\n23.75\nDeep-Att + PosUnk [39]\n39.2\n1.0 · 1020\nGNMT + RL [38]\n24.6\n39.92\n2.3 · 1019\n1.4 · 1020\nConvS2S [9]\n25.16\n40.46\n9.6 · 1018\n1.5 · 1020\nMoE [32]\n26.03\n40.56\n2.0 · 1019\n1.2 · 1020\nDeep-Att + PosUnk Ensemble [39]\n40.4\n8.0 · 1020\n\n\n
    '''
    prompt2 = '你好，怎么做西红柿炒鸡蛋？'


    llama3_8b_config = LLAMA3_CONFIG(engine="meta/llama3-8b-instruct")
    llama3_70b_config = LLAMA3_CONFIG(engine="meta/llama3-70b-instruct")
    llama3_instance = LLAMA3(llama3_70b_config)
    # 给出prompt，调用API
    result, _ = llama3_instance.get_instructions(prompt_chinese)
    print(result)
    print(_.json())
